from absl import app

import collections
import os
import logging
import glob
import re


def extract_pattern(filename, regs):
  pattern_chosen = [
      fnpart for fnpart in filename.split(".") if any([reg.match(fnpart) for reg in regs])
  ]
  assert len(pattern_chosen) == 1, (filename, regs)
  return pattern_chosen[0]


def extract_edge(filename):
  stk1 = re.compile(".*[0-9]*bp")
  stk2 = re.compile("e[0-9]*")
  return extract_pattern(filename, [stk1, stk2])


def extract_stack(filename):
  stk1 = re.compile("^[0-9]*stk")
  stk2 = re.compile("^s[0-9]*")
  return extract_pattern(filename, [stk1, stk2])


def main(_):
  pics = glob.glob("oe_stat_backtest/pnl_fill.*")
  picgroups = set([pic.split(".")[1] for pic in pics])
  for picgroup in picgroups:
    picfiles = glob.glob("oe_stat_backtest/pnl_fill.%s.*" % picgroup)
    allstacks = sorted(list(set([extract_stack(picfile) for picfile in picfiles])))
    alledges = sorted(list(set([extract_edge(picfile) for picfile in picfiles])))
    for legacy in [False, True]:
      if legacy:
        stacks = [stack for stack in allstacks if stack.endswith("stk")]
        edges = [edge for edge in alledges if edge.endswith("bp")]
      else:
        stacks = [stack for stack in allstacks if stack.startswith("s")]
        edges = [edge for edge in alledges if edge.startswith("e")]
      if len(stacks) == 0:
        continue
      logging.info("group %s stks %s edgs %s", picgroup, stacks, edges)

      picgroupfile = "oe_stat_backtest/pnl_combined_%s_%s_%s.png" % (
          picgroup, "_".join(stacks), "_".join(edges))

      while True:
        allpics = []
        subcmds = []

        stackscnt = collections.defaultdict(int)
        edgescnt = collections.defaultdict(int)

        for stack in stacks:
          pfs_stack = [picfile for picfile in picfiles if picfile.find(stack) >= 0]
          pfchosens = []
          for edge in edges:
            pfchosen = [picfile for picfile in pfs_stack if picfile.find(edge) >= 0]
            # assert len(pfchosen) == 1, (stack, edge, pfchosen, pfs_stack, picfiles)
            if len(pfchosen) > 0:
              pfchosens.append(pfchosen[0])
              allpics.append(pfchosen[0])
              stackscnt[stack] += 1
              edgescnt[edge] += 1
          subcmd = "( %s +append ) \\" % (" ".join(pfchosens))
          subcmds.append(subcmd)
        if len(set(edgescnt.values())) > 1:
          max_edge_cnt = max(set(edgescnt.values()))
          edges = [edge for edge, cnt in edgescnt.items() if cnt == max_edge_cnt]
        else:
          break
      # print(picgroup, stackscnt, edgescnt)
      # import pdb; pdb.set_trace()
      cmd_mtg = "montage -mode concatenate -tile %sx%s %s %s" % (
          len(edges), len(stacks), " ".join(allpics), picgroupfile)
      print(cmd_mtg)
      os.system(cmd_mtg)

      # cmd = "convert %s \n -append %s" % ("\n".join(subcmds), picgroupfile)
      # print(cmd)
      # os.system(cmd)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
